import re

import numpy

from pydap.model import *


def mean(var, axis=0):
    """
    Returns the average along a given axis.

    :param var: an Opendap variable
    :type var: Sequence/basic type

    :param axis: axis along which the average will be performed (0-indexed)
    :type axis: integer

    :return: an Opendap variable
    :rtype: Sequence/basic type

    """
    if isinstance(var, BaseType):
        var.data = numpy.mean(var, axis=axis)
        var.shape = var.data.shape
        dims = list(var.dimensions)
        dims.pop(int(axis))
        var.dimensions = tuple(dims)
    elif isinstance(var, SequenceType):
        var.data = [ numpy.mean(list(var.data), axis) ]
    return var


def grid(var, lat1, lon0, lat0, lon1, time=None): 
    slice_ = [slice(None)] * len(var.shape)

    lat = get_lat(var)
    s = (lat0 < numpy.asarray(lat)) & (numpy.asarray(lat) < lat1)
    start = s.nonzero()[0][0]
    stop = s.nonzero()[0][-1] + 1
    slice_[list(var.dimensions).index(lat.name)] = slice(start, stop, 1)

    lon = get_lon(var)
    s = (lon0 < numpy.asarray(lon)) & (numpy.asarray(lon) < lon1)
    start = s.nonzero()[0][0]
    stop = s.nonzero()[0][-1] + 1
    slice_[list(var.dimensions).index(lon.name)] = slice(start, stop, 1)

    if time is not None:
        pass

    return var[ tuple(slice_) ]
    

def geoarray(var): pass
def geogrid(var): pass


def linear_scale(var, m, b): 
    """
    Scale dat using ``y=mx+b``.

    :param var: an Opendap variable
    :type var: Sequence/basic type

    :param m: multiplier
    :type axis: float

    :param b: offset
    :type axis: float

    """
    if isinstance(var, BaseType):
        var.data = numpy.asarray(var.data)*m+b
    return var


def get_lon(grid):
    lon = [dim for dim in grid.maps.values() if
            re.match('degrees?_e', dim.attributes.get('units', ''))]
    if lon:
        return lon[0]
    else:
        return None


def get_lat(grid):
    lat = [dim for dim in grid.maps.values() if
            re.match('degrees?_n', dim.attributes.get('units', ''))]
    if lat:
        return lat[0]
    else:
        return None
